#!/usr/bin/python


## Alvar marker extraction script
## Created Jan 28, 2011
## Author: Sean Bowman, University of Minnesota

import sys
import os
import xml.dom.minidom
import re

# suppress alvar's writing to stderr when imported... maybe this is overkill.
#devnull = open(os.devnull, 'w')
#save_stderr = os.dup(2)
#os.dup2(devnull.fileno(), 2)
import alvar
#os.dup2(save_stderr, 2)

## parse the XML file containing real-world coordinate information
xml_doc = xml.dom.minidom.parse("MarkerData.xml")

def usage():
	print 
	print "This script scans all images that fit the format "
	print "\t[base_file_name]-[#].[extension]"
	print "where base_file_name and extension can be specified, and # is any number."
	print
	print "Usage: %s [options]" % sys.argv[0] 
	print
	print "Available options:"
	print "\t --ext=<type> \t\t the input image extension [pgm]"
	print "\t --indir=<dir> \t\t directory containing the input image files [.]"
	print "\t --outdir=<dir> \t\t the directory to write the output files [.]"
	print "\t --filename=<name> \t the base file name [image]"
	print "\t --marker-limit=<num> \t the max number of detectable markers in one image [50]"
	print 
	
if __name__ == '__main__':
	if len(sys.argv) >= 2 and sys.argv[1] == '--help':
		usage(); exit(0)
	
	# set default values
	inputdir = '.'
	extension = 'pgm'
	outputdir = '.'
	base_file_name = "image"
	max_markers = 50
	
	# parse any present command line options
	for arg in sys.argv[1:]:
		flag, eq, val = arg.partition('=')
		if flag == '--ext':
			extension = val
		elif flag == '--indir':
			# strip a trailing '/' character if it is present
			if val[-1] == '/':
				val = val[:-1]
			inputdir = val
		elif flag == '--outdir':
			# strip a trailing '/' character if it is present
			if val[-1] == '/':
				val = val[:-1]
			outputdir = val
		elif flag == '--filename':
			base_file_name = val
		elif flag == '--marker-limit':
			max_markers = int(val)
		else:
			print
			print "Invalid argument: %s" % arg
			print "Use --help for more info."
			exit(1)
			
	alvar.init(max_markers) # TODO fix the memory leaks here!!
	
	
	# Some detail on the "base" name of the input files:
	# 	Files that fit the format base_file_name-num.[extension] will be analyzed
	# 	by the script, where num is any number. Example: with base_file_name = "im"
	# 	and extension = jpg, some valid images are im-3.jpg, im-01392.jpg, etc.
	
		
	# The input and output file IDs will not always match - for example, if we have
	# input files "img-1.jpg" and "img-3.jpg", their outputs will still be "x_1" and 
	# "x_2". This correspondences file will record the mapping between the two.
	# It contains two columns: the left is the image file ID, the right is the output ID.
	correspondence_file = open(outputdir + "/file_correspondences.dat", 'w')
	
	# iterate over all valid files in the input directory
	# uses a simple regex to match file names exactly
	files = (f for f in sorted(os.listdir(inputdir)) if re.match(base_file_name + "-[0-9]+." + extension, f) != None)
	num_written = 0
	output_id = 0
	for file in files:
		num_markers = alvar.get_markers(inputdir + '/' + file)
			
		## keep a set of detected marker IDs so we don't get the same one twice,
		## as well as a set of "good" markers (ones detected only once) so we 
		## know which to write
		detected_ids = set()
		good_marker_ids = set()
		
		## iterate over all detected markers. we need to do this twice; once
		## to check for which markers are not duplicated, and once to write
		## those markers to the output files.
		for i in range(num_markers):
			marker_id = alvar.get_id(i)
			if marker_id in detected_ids:
				# duplicate!
				print "Marker %d is duplicated in %s, ignoring it" \
						% (marker_id, file)
				if marker_id in good_marker_ids:
					good_marker_ids.remove(marker_id)
			else:
				detected_ids.add(marker_id)
				good_marker_ids.add(marker_id)
		
		# make sure this image has detectable markers -- we don't want to
		# write a blank file.
		if len(good_marker_ids) == 0:
			continue
			
		# this file is good, begin writing data
		output_id += 1
		## get the file ID, and record the correspondence
		file_id = re.search('[0-9]+', file).group(0)
		correspondence_file.write(file_id + " " + str(output_id) + "\n")
			
		# open output files
		image_coords_out = open(outputdir + "/x_" + str(output_id), 'w')
		real_coords_out = open(outputdir + "/X_" + str(output_id), 'w')
		
		## iterate over all detected markers again to write their data
		for j in range(num_markers):
			marker_id = alvar.get_id(j)
			
			# check that this isn't one of the duplicated markers
			if marker_id not in good_marker_ids:
				continue
			
			## get the XML entry corresponding to this marker
			marker_xml = [m for m in xml_doc.getElementsByTagName("marker") if int(m.getAttribute("index")) == marker_id]
			if len(marker_xml) == 0:
				print "Error: marker ID not found in XML file."
				print "Continuing to next marker..."
			else:
				## iterate over all corners of this marker, printing both their
				## image and real coordinates to the output files
				corners = marker_xml[0].getElementsByTagName("corner")
				for k in range(4):
					image_data = str(alvar.get_corners(j, k).x) + " " + \
								 str(alvar.get_corners(j, k).y)
					image_coords_out.write(image_data + "\n")
					
					real_data = corners[k].getAttribute("x") + " " + \
								corners[k].getAttribute("y") + " " + \
								corners[k].getAttribute("z")
					real_coords_out.write(real_data + "\n")
			
					
		image_coords_out.close()
		real_coords_out.close()
		num_written = num_written + 1
	
	print
	if num_written > 0:
		print "Successfully wrote data to %d files!" % num_written
	else:
		print "Found 0 images. Did you specify the correct options?"
		print "Use --help for more info."
		print
